"""
卡尔曼滤波策略(zipline实现)
策略逻辑：
1. 使用卡尔曼滤波估计价格均衡水平
2. 当价格偏离均衡水平超过阈值时交易
"""

from zipline.api import order_target, record, symbol
import numpy as np
from pykalman import KalmanFilter

def initialize(context):
    context.asset = symbol('AAPL')
    context.threshold = 1.0  # 交易阈值
    
    # 卡尔曼滤波参数
    context.transition_covariance = 0.01
    context.observation_covariance = 0.1

def handle_data(context, data):
    # 获取历史价格
    prices = data.history(context.asset, 'price', 100, '1d')
    
    # 初始化卡尔曼滤波
    kf = KalmanFilter(
        transition_matrices=[1],
        observation_matrices=[1],
        initial_state_mean=prices[0],
        initial_state_covariance=1,
        transition_covariance=context.transition_covariance,
        observation_covariance=context.observation_covariance
    )
    
    # 运行卡尔曼滤波
    state_means, _ = kf.filter(prices.values)
    predicted = state_means[-1][0]
    
    # 计算偏离程度
    current_price = prices[-1]
    deviation = (current_price - predicted) / predicted
    
    # 交易逻辑
    current_position = context.portfolio.positions[context.asset].amount
    
    if deviation < -context.threshold and current_position <= 0:
        order_target(context.asset, 100)  # 买入
    elif deviation > context.threshold and current_position > 0:
        order_target(context.asset, 0)  # 卖出
    
    # 记录预测值和偏离度
    record(predicted=predicted, deviation=deviation)

if __name__ == '__main__':
    from zipline.utils.run_algo import run_algorithm
    # 这里添加回测配置